package org.luxor.commons.selenium.component;

import cn.hutool.http.HtmlUtil;
import com.google.common.collect.ImmutableMap;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.luxor.commons.selenium.component.storage.NoneLocalStorage;
import org.luxor.commons.selenium.component.storage.NoneSessionStorage;
import org.luxor.commons.selenium.config.properties.SeleniumProperties;
import org.luxor.commons.selenium.exception.StealthJsException;
import org.openqa.selenium.*;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chromium.HasCdp;
import org.openqa.selenium.devtools.DevTools;
import org.openqa.selenium.devtools.DevToolsException;
import org.openqa.selenium.devtools.HasDevTools;
import org.openqa.selenium.devtools.v112.network.Network;
import org.openqa.selenium.devtools.v112.network.model.RequestId;
import org.openqa.selenium.html5.LocalStorage;
import org.openqa.selenium.html5.SessionStorage;
import org.openqa.selenium.html5.WebStorage;
import org.openqa.selenium.remote.Augmenter;
import org.openqa.selenium.remote.ExecuteMethod;
import org.openqa.selenium.remote.RemoteExecuteMethod;
import org.openqa.selenium.remote.RemoteWebDriver;
import org.openqa.selenium.support.events.EventFiringWebDriver;
import org.openqa.selenium.support.ui.ExpectedCondition;
import org.openqa.selenium.support.ui.ExpectedConditions;
import org.openqa.selenium.support.ui.WebDriverWait;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.CollectionUtils;
import org.springframework.util.PathMatcher;

import java.io.File;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

import static org.openqa.selenium.remote.CapabilityType.PLATFORM;
import static org.openqa.selenium.remote.CapabilityType.PLATFORM_NAME;

/**
 * 基于扩展EventFiringWebDriver进行扩展，实现定制化能力
 *
 * @author Mr.yan  @date 2022/5/17
 */
public class SeleniumDriver extends EventFiringWebDriver implements WebStorage, HasCdp, HasDevTools {
    private static final Logger log = LoggerFactory.getLogger(SeleniumDriver.class);

    private static final PathMatcher pathMatcher = new AntPathMatcher();

    private final List<String> supportWebStorageDrivers = Arrays.asList("ChromeDriver", "FirefoxDriver", "OperaDriver");

    private final NoneLocalStorage noneLocalStorage = new NoneLocalStorage();

    private final NoneSessionStorage noneSessionStorage = new NoneSessionStorage();

    private final SeleniumProperties properties;

    private final RemoteWebDriver webDriver;

    private final Process process;

    private final String userDataDir;

    public SeleniumDriver(RemoteWebDriver webDriver, SeleniumProperties properties, Process process, String userDataDir) {
        super(webDriver);
        this.webDriver = webDriver;
        this.properties = properties;
        this.process = process;
        this.userDataDir = userDataDir;
    }

    @Override
    public void quit() {
        webDriver.quit();
        if (process != null) {
            process.destroyForcibly();
        }
        if (userDataDir != null) {
            FileUtils.deleteQuietly(new File(userDataDir));
        }
    }

    public RemoteWebDriver webDriver() {
        return this.webDriver;
    }

    public String getSessionId() {
        return this.webDriver.getSessionId().toString();
    }


    @Override
    public WebElement findElement(By by) {
        return super.findElement(by);
    }

    public WebElement findElement(By by, long timeoutSeconds) {
        try {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(timeoutSeconds));
            return super.findElement(by);
        } finally {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(properties.getImplicitlyWaitTimeOutSeconds()));
        }
    }

    public synchronized WebElement findClickElement(By by) {
        return new WebDriverWait(webDriver, Duration.ofSeconds(properties.getImplicitlyWaitTimeOutSeconds())).until(ExpectedConditions.elementToBeClickable(by));
    }

    public synchronized WebElement findClickElement(By by, long timeoutSeconds) {
        try {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(timeoutSeconds));
            return new WebDriverWait(webDriver, Duration.ofSeconds(timeoutSeconds)).until(ExpectedConditions.elementToBeClickable(by));
        } finally {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(properties.getImplicitlyWaitTimeOutSeconds()));
        }
    }

    @Override
    public List<WebElement> findElements(By by) {
        return super.findElements(by);
    }

    public List<WebElement> findElements(By by, long timeoutSeconds) {
        try {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(timeoutSeconds));
            return super.findElements(by);
        } finally {
            manage().timeouts().implicitlyWait(Duration.ofSeconds(properties.getImplicitlyWaitTimeOutSeconds()));
        }
    }

    public List<WebElement> findClickElements(By by) {
        List<WebElement> clickElements = new ArrayList<>();
        List<WebElement> elementList = this.findElements(by);
        for (WebElement element : elementList) {
            WebElement clickElement = new WebDriverWait(webDriver, Duration.ofSeconds(1000)).until(ExpectedConditions.elementToBeClickable(element));
            clickElements.add(clickElement);
        }
        return clickElements;
    }

    public List<WebElement> findClickElements(By by, long timeoutSeconds) {
        List<WebElement> clickElements = new ArrayList<>();
        List<WebElement> elementList = this.findElements(by, timeoutSeconds);
        for (WebElement element : elementList) {
            WebElement clickElement = new WebDriverWait(webDriver, Duration.ofSeconds(1000)).until(ExpectedConditions.elementToBeClickable(element));
            clickElements.add(clickElement);
        }
        return clickElements;
    }

    public Screenshot createScreenshot() {
        Screenshot screenshot = new Screenshot();
        try {
            screenshot.setHtmlText(HtmlUtil.removeHtmlTag(webDriver().getPageSource(), "path", "symbol", "style"));
        } catch (Throwable e) {
            // .ignore
        }
        try {
            screenshot.setBase64Image(webDriver().getScreenshotAs(OutputType.BASE64));
        } catch (Throwable e) {
            // .ignore
        }
        return screenshot;
    }

    @Override
    public LocalStorage getLocalStorage() {
        String webDriveName = webDriver().getClass().getSimpleName();
        if (supportWebStorageDrivers.contains(webDriveName)) {
            return ((WebStorage) webDriver()).getLocalStorage();
        }
        log.warn("Used to indicate that a localStorage used by the {} is unsupported.", webDriveName);
        return this.noneLocalStorage;
    }

    @Override
    public SessionStorage getSessionStorage() {
        String webDriveName = webDriver().getClass().getSimpleName();
        if (supportWebStorageDrivers.contains(webDriveName)) {
            return ((WebStorage) webDriver()).getSessionStorage();
        }
        log.warn("Used to indicate that a sessionStorage used by the {} is unsupported.", webDriveName);
        return this.noneSessionStorage;
    }

    /**
     * 重置浏览器cookies
     *
     * @param baseUrl 地址 (建议使用稳定且响应时间最短的资源)
     * @param cookies
     */
    public synchronized void resetCookies(String baseUrl, List<Cookie> cookies) {
        if (CollectionUtils.isEmpty(cookies)) {
            return;
        }
        this.webDriver().get(baseUrl);
        this.webDriver().manage().deleteAllCookies();
        for (Cookie cookie : cookies) {
            this.webDriver().manage().addCookie(cookie);
        }
    }

    /**
     * 执行爬虫特征伪装脚本
     *
     * @param autoRefresh 执行后自动刷新吗? true-是；false-否；
     */
    public synchronized void stealthJs(boolean autoRefresh) {
        try {
            ClassPathResource stealthJsResource = new ClassPathResource("org/luxor/commons/selenium/js/stealth.min.js");
            String stealthJs = IOUtils.toString(stealthJsResource.getInputStream(), StandardCharsets.UTF_8);
            Map<String, Object> uaMap = new HashMap<String, Object>(2) {{
                put("source", stealthJs);
            }};
            this.executeCdpCommand("Page.addScriptToEvaluateOnNewDocument", uaMap);
        } catch (Exception e) {
            throw new StealthJsException("爬虫特征码抹除脚本(stealth.min.js),执行失败", e);
        }
        try {
            ClassPathResource stealthExtendJsResource = new ClassPathResource("org/luxor/commons/selenium/js/stealth.extend.js");
            String stealthExtendJs = IOUtils.toString(stealthExtendJsResource.getInputStream(), StandardCharsets.UTF_8);
            Map<String, Object> uaMap = new HashMap<String, Object>(2) {{
                put("source", stealthExtendJs);
            }};
            this.executeCdpCommand("Page.addScriptToEvaluateOnNewDocument", uaMap);
        } catch (Exception e) {
            throw new StealthJsException("爬虫特征码抹除脚本(stealth.extend.js),执行失败", e);
        }
        if (autoRefresh) {
            this.webDriver().navigate().refresh();
        }
    }

    @Override
    public Map<String, Object> executeCdpCommand(String commandName, Map<String, Object> parameters) {
        Objects.requireNonNull(commandName, "Command name must be set.");
        Objects.requireNonNull(parameters, "Parameters for command must be set.");

        ExecuteMethod executeMethod = new RemoteExecuteMethod(this.webDriver());
        Map<String, Object> toReturn = (Map) executeMethod.execute("executeCdpCommand", ImmutableMap.of("cmd", commandName, "params", parameters));
        return ImmutableMap.copyOf(toReturn);
    }

    @Override
    public Optional<DevTools> maybeGetDevTools() {
        if (this.webDriver() instanceof ChromeDriver) {
            return ((HasDevTools) this.webDriver()).maybeGetDevTools();
        } else if (this.webDriver() instanceof RemoteWebDriver) {
            return Optional.ofNullable(((HasDevTools) new Augmenter().augment(this.webDriver())).getDevTools());
        }
        return Optional.empty();
    }

    @Override
    public DevTools getDevTools() {
        Optional<DevTools> devTools = maybeGetDevTools();
        if (devTools.isPresent()) {
            devTools.get().createSessionIfThereIsNotOne();
            devTools.get().send(Network.enable(Optional.empty(), Optional.empty(), Optional.empty()));
        }
        return devTools.orElseThrow(() -> new DevToolsException("Unable to create DevTools connection"));
    }

    public AtomicReference<RequestId> addRequestListener(String urlPattern) {
        AtomicReference<RequestId> request = new AtomicReference<>();
        this.getDevTools().addListener(Network.requestWillBeSent(), req -> {
            if (pathMatcher.match(urlPattern, req.getRequest().getUrl())) {
                request.getAndSet(req.getRequestId());
                log.info("Condition [{}] matches the RequestId:{}, url {}", urlPattern, req.getRequestId(), req.getRequest().getUrl());
            }
        });
        return request;
    }

    public boolean waitForRequest(AtomicReference<RequestId> requestId, Duration timeout) {
        // wait requestId
        new WebDriverWait(webDriver, timeout, Duration.ofMillis(500)).until((ExpectedCondition<Boolean>) driver -> requestId.get() != null);
        return true;
    }

    public AtomicReference<RequestId> addResponseListener(String urlPattern) {
        AtomicReference<RequestId> requestId = new AtomicReference<>();
        this.getDevTools().addListener(Network.responseReceived(), res -> {
            if (pathMatcher.match(urlPattern, res.getResponse().getUrl())) {
                requestId.getAndSet(res.getRequestId());
                log.info("Condition [{}] matches the RequestId:{}, url {}", urlPattern, res.getRequestId(), res.getResponse().getUrl());
            }
        });
        return requestId;
    }

    public String waitForResponseBody(AtomicReference<RequestId> requestId, Duration timeout) {
        Optional<DevTools> devTools = maybeGetDevTools();
        if (!devTools.isPresent()) {
            throw new DevToolsException("Unable to get DevTools connection");
        }
        try {
            // wait requestId
            new WebDriverWait(webDriver, timeout, Duration.ofMillis(500)).until((ExpectedCondition<Boolean>) driver -> requestId.get() != null);
            return devTools.get().send(Network.getResponseBody(requestId.get())).getBody();
        } catch (TimeoutException e) {
            throw e;
        }
    }

    @Override
    public String toString() {
        Capabilities caps = getCapabilities();
        if (caps == null) {
            return super.toString();
        }

        // w3c name first
        Object platform = caps.getCapability(PLATFORM_NAME);
        if (!(platform instanceof String)) {
            platform = caps.getCapability(PLATFORM);
        }
        if (platform == null) {
            platform = "unknown";
        }

        return String.format("SeleniumDriver(%s): %s on %s (%s)", webDriver().getClass().getSimpleName(), caps.getBrowserName(), platform, getSessionId());
    }
}
